import tensorflow as tf
import tensorflow_addons as tfa

class gatedLinearLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(gatedLinearLayer,self).__init__()
    def call(self, inputs, gate):
        activation = tf.multiply(x = inputs, y = tf.sigmoid(gate))
        return activation

class residual1dBlock(tf.keras.layers.Layer):
    def __init__(self,
                 filters=1024,
                 kernel_size=3,
                 strides=1,
                 padding="SAME"
                 ):
        super(residual1dBlock, self).__init__()
        # inputs = tf.keras.Input()
        self.conv1 = tf.keras.layers.Conv1D(filters = filters,
                                            kernel_size = kernel_size,
                                            strides = strides,
                                            padding = padding)
        self.conv1Norm = tfa.layers.InstanceNormalization()
        self.cnnGate = tf.keras.layers.Conv1D(filters = filters,
                                              kernel_size = kernel_size,
                                              strides = strides,
                                              padding = padding)
        self.cnnGateNorm = tfa.layers.InstanceNormalization()
        self.glu = gatedLinearLayer()
        self.conv2 = tf.keras.layers.Conv1D(filters = filters // 2,
                                            kernel_size = kernel_size,
                                            strides = strides,
                                            padding = padding)
        self.conv2Norm = tfa.layers.InstanceNormalization()

    def call(self, inputs, **kwargs):
        conv1 = self.conv1(inputs)
        conv1Norm = self.conv1Norm(conv1)
        cnnGate = self.cnnGate(inputs)
        cnnGateNorm = tfa.layers.InstanceNormalization()(cnnGate)
        glu = self.glu(conv1Norm, cnnGateNorm)
        conv2 =self.conv2(glu)
        conv2Norm = self.conv2Norm(conv2)
        return conv2Norm + inputs

class downSample1DBlock(tf.keras.layers.Layer):
    def __init__(self,
                 filters,
                 kernel_size,
                 strides,
                 padding="SAME"
                 ):
        super(downSample1DBlock, self).__init__()

        self.conv1 = tf.keras.layers.Conv1D(filters = filters,
                                       kernel_size = kernel_size,
                                       strides = strides,
                                       padding = padding)
        self.conv1Norm = tfa.layers.InstanceNormalization()
        self.cnnGate = tf.keras.layers.Conv1D(filters = filters,
                                         kernel_size = kernel_size,
                                         strides = strides,
                                         padding = padding)
        self.cnnGateNorm = tfa.layers.InstanceNormalization()
        self.glu = gatedLinearLayer()

    def call(self, inputs, **kwargs):
        conv1 = self.conv1(inputs)
        conv1Norm = self.conv1Norm(conv1)
        cnnGate = self.cnnGate(inputs)
        cnnGateNorm = self.cnnGateNorm(cnnGate)
        glu = self.glu(conv1Norm, cnnGateNorm)
        return glu

class downSample2DBlock(tf.keras.layers.Layer):
    def __init__(self,
                 filters,
                 kernel_size,
                 strides,
                 padding="SAME"
                 ):
        super(downSample2DBlock, self).__init__()

        self.conv1 = tf.keras.layers.Conv2D(filters=filters,
                                            kernel_size=kernel_size,
                                            strides=strides,
                                            padding=padding)
        self.conv1Norm = tfa.layers.InstanceNormalization()
        self.cnnGate = tf.keras.layers.Conv2D(filters=filters,
                                              kernel_size=kernel_size,
                                              strides=strides,
                                              padding=padding)
        self.cnnGateNorm = tfa.layers.InstanceNormalization()
        self.glu = gatedLinearLayer()

    def call(self, inputs, **kwargs):
        conv1 = self.conv1(inputs)
        conv1Norm = self.conv1Norm(conv1)
        cnnGate = self.cnnGate(inputs)
        cnnGateNorm = self.cnnGateNorm(cnnGate)
        glu = self.glu(conv1Norm, cnnGateNorm)
        return glu


# def downSample2DBlock(inputs,
#                       filters,
#                       kernel_size,
#                       strides,
#                       padding = "SAME"):
#     conv1 = tf.keras.layers.Conv2D(filters=filters,
#                                    kernel_size=kernel_size,
#                                    strides=strides,
#                                    padding=padding,
#                                    )(inputs)
#     conv1Norm = tfa.layers.InstanceNormalization()(conv1)
#     cnnGate = tf.keras.layers.Conv2D(filters=filters,
#                                      kernel_size=kernel_size,
#                                      strides=strides,
#                                      padding=padding)(inputs)
#     cnnGateNorm = tfa.layers.InstanceNormalization()(cnnGate)
#     glu = gatedLinearLayer(conv1Norm, cnnGateNorm)
#     return glu

class upSample1DBlock(tf.keras.layers.Layer):
    def __init__(self,
                 filters,
                 kernel_size,
                 strides,
                 shuffle_size=2,
                 padding="SAME"
                 ):
        super(upSample1DBlock, self).__init__()

        self.conv1 = tf.keras.layers.Conv1D(filters = filters,
                                            kernel_size = kernel_size,
                                            strides = strides,
                                            padding = padding)
        self.conv1Shuffle = pixelShuffle(shuffle_size = shuffle_size)
        self.conv1Norm = tfa.layers.InstanceNormalization()

        self.cnnGate = tf.keras.layers.Conv1D(filters = filters,
                                              kernel_size = kernel_size,
                                              strides = strides,
                                              padding = padding)
        self.cnnGateShuffle = pixelShuffle(shuffle_size = shuffle_size)
        self.cnnGateNorm = tfa.layers.InstanceNormalization()
        self.glu = gatedLinearLayer()

    def call(self, inputs, **kwargs):
        conv1 = self.conv1(inputs)
        conv1Shuffle = self.conv1Shuffle(conv1)
        conv1Norm = self.conv1Norm(conv1Shuffle)

        cnnGate = self.cnnGate(inputs)
        cnnGateShuffle = self.cnnGateShuffle(cnnGate)
        cnnGateNorm = self.cnnGateNorm(cnnGateShuffle)

        glu = self.glu(conv1Norm, cnnGateNorm)
        return glu

class pixelShuffle(tf.keras.layers.Layer):
    def __init__(self, shuffle_size = 2):
        super(pixelShuffle, self).__init__()
        self.shuffle_size = shuffle_size

    def call(self, inputs, **kwargs):
        h = tf.shape(inputs)[0]
        w = tf.shape(inputs)[1]
        c = inputs.get_shape().as_list()[2]
        outputC = c // self.shuffle_size
        outputW = w * self.shuffle_size
        outputs = tf.reshape(tensor = inputs,
                             shape = [h, outputW, outputC])
        return outputs


def generatorWithGatedCNN():
    inputs1 = tf.keras.Input(shape=(24, 128))
    # inputShape = tf.shape(inputs1)
    inputs = tf.keras.layers.Reshape(target_shape=(128, 24))(inputs1)
    conv1 = tf.keras.layers.Conv1D(filters = 128,
                                   kernel_size = 15,
                                   strides = 1,
                                   padding="same"
                                   )(inputs)
    conv1Gate = tf.keras.layers.Conv1D(filters = 128,
                                       kernel_size=15,
                                       strides=1,
                                       padding="same")(inputs)
    glu1 = gatedLinearLayer()(conv1, conv1Gate)
    d1 = downSample1DBlock(filters=256,
                           kernel_size=5,
                           strides=2)(glu1)
    d2 = downSample1DBlock(filters=512,
                           kernel_size=5,
                           strides=2)(d1)
    r1 = residual1dBlock(filters=1024,
                         kernel_size=3,
                         strides=1)(d2)
    r2 = residual1dBlock(filters=1024,
                         kernel_size=3,
                         strides=1)(r1)
    r3 = residual1dBlock(filters=1024,
                         kernel_size=3,
                         strides=1)(r2)
    r4 = residual1dBlock(filters=1024,
                         kernel_size=3,
                         strides=1)(r3)
    r5 = residual1dBlock(filters=1024,
                         kernel_size=3,
                         strides=1)(r4)
    r6 = residual1dBlock(filters=1024,
                         kernel_size=3,
                         strides=1)(r5)
    u1 = upSample1DBlock(filters=1024,
                         kernel_size=5,
                         strides=1,
                         shuffle_size=2)(r6)
    u2 = upSample1DBlock(filters=512,
                         kernel_size=5,
                         strides=1,
                         shuffle_size=2)(u1)
    o1 = tf.keras.layers.Conv1D(filters=24,
                                kernel_size=15,
                                strides=1,
                                padding = "SAME")(u2)
    o2 = tf.transpose(o1, perm = [0, 2, 1])
    model = tf.keras.Model(inputs = inputs1, outputs = o2)
    return model

def disciriminator():
    inputs1 = tf.keras.layers.Input(shape=(24,128))
    inputs = tf.keras.layers.Reshape(target_shape=[24, 128, 1])(inputs1)
    conv1 = tf.keras.layers.Conv2D(filters=128,
                                   kernel_size=[3,3],
                                   strides=[1,2],
                                   padding="SAME")(inputs)
    cnnGate = tf.keras.layers.Conv2D(filters=128,
                                     kernel_size=[3,3],
                                     strides=[1,2],
                                     padding="SAME")(inputs)
    glu1 = gatedLinearLayer()(conv1, cnnGate)
    d1 = downSample2DBlock(filters=256,
                           kernel_size=[3,3],
                           strides=[2,2],
                           padding="SAME")(glu1)
    d2 = downSample2DBlock(filters=512,
                           kernel_size=[3, 3],
                           strides=[2, 2],
                           padding="SAME")(d1)
    d3 = downSample2DBlock(filters=1024,
                           kernel_size=[6, 3],
                           strides=[1, 2],
                           padding="SAME")(d2)
    fl = tf.keras.layers.Flatten()(d3)
    o1 = tf.keras.layers.Dense(units=1, activation="sigmoid")(fl)
    model = tf.keras.Model(inputs=inputs1, outputs=o1)
    return model

g = generatorWithGatedCNN()
d = disciriminator()
print(d.summary())
print(g.summary())
